from keras import Input
from util.model_listutil import ModelListUtil
from configuration.model_congig import ModelConfig
from datamanager import data_manager
from model import Model
from plot import Plot

labels = [
    "Harmful",
    "Kitchen",
    "Other",
    "Recyclable",
    "micro",
    "macro",
]


class ReDraw:
    def __init__(self, image_height, image_width):
        self.conf = ModelConfig()
        self.conf.IMAGE_HEIGHT = image_height
        self.conf.IMAGE_WIDTH = image_width
        self.inputs = Input(shape=(self.conf.IMAGE_HEIGHT, self.conf.IMAGE_HEIGHT, self.conf.COLOR_CHANNELS))

    def draw(self, train_id: int):
        mtest = Model(self.conf, self.inputs)
        model_list = ModelListUtil.get_model_list_from_db(train_id)
        mtest.load(model_list, f"storage/{train_id}.h5")
        roc_img_name = f"ROC{train_id}.png"
        pr_img_name = f"PR{train_id}.png"
        y_pred = mtest.predict(data_manager.x_test)
        Plot().plot_ROC(data_manager.y_test, y_pred, labels, roc_img_name)
        Plot().plot_pr(data_manager.y_test, y_pred, labels, pr_img_name)


if __name__ == '__main__':

    # ReDraw().draw(332)
    # ReDraw().draw(333)
    # ReDraw().draw(336)
    # ReDraw().draw(335)
    # ReDraw().draw(367)
    # ReDraw(128,128).draw(483)
    # for i in range(337, 550):
    #     print(i)
    #     # try:
    #     #     ReDraw(128, 128).draw(i)
    #     # except Exception as e:
    #     #     print(e)
    #     #     continue
    #     # try:
    #     #     ReDraw(96, 96).draw(i)
    #     # except Exception as e:
    #     #     print(e)
    #     #     continue
    #     try:
    #         ReDraw(128, 128).draw(i)
    #     except Exception as e:
    #         print(e)
    #         continue
    ReDraw(128, 128).draw(502)
